import cv2
import numpy as np
from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import StandardScaler  
import os  

cat_paths = [f"./dataset/Cat/{i}.jpg" for i in range(0, 101)]  
dog_paths = [f"./dataset/Dog/{i}.jpg" for i in range(0, 101)]  
TEST_CAT_DIR = r"./dataset/test/Cat"
TEST_DOG_DIR = r"./dataset/test/Dog"
IMAGE_EXTENSIONS = (".jpg", ".jpeg", ".png", ".bmp", ".gif")

def extract_manual_features(img_path):
    img = cv2.imread(img_path)
    if img is None:
        raise ValueError(f"无法读取图像：{img_path}，请检查路径是否正确或文件是否损坏")
    H, W, _ = img.shape

    #变量R、G、B读取通道
    B, G, R = cv2.split(img)
    #R, G, B = img[::0], img[::1], img[::2]
    
    # 特征1：平均亮度（手工计算灰度图）
    gray_manual = 0.299 * R + 0.587 * G + 0.114 * B
    mean_brightness = np.mean(gray_manual)
    

    sobel_x = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
    sobel_y = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])
    edges_x, edges_y = np.zeros_like(gray_manual), np.zeros_like(gray_manual)
    for i in range(1, H-1):
        for j in range(1, W-1):
            edges_x[i, j] = np.sum(gray_manual[i-1:i+2, j-1:j+2] * sobel_x)
            edges_y[i, j] = np.sum(gray_manual[i-1:i+2, j-1:j+2] * sobel_y)
    edge_strength = np.sqrt(edges_x**2 + edges_y**2)
    edge_density = np.sum(edge_strength > 50) / (H * W)  
    
    # 特征3：颜色饱和度（手工计算）
    R_norm, G_norm, B_norm = R / 255.0, G / 255.0, B / 255.0
    min_rgb = np.minimum(np.minimum(R_norm, G_norm), B_norm)
    max_rgb = np.maximum(np.maximum(R_norm, G_norm), B_norm)
    saturation_pixel = np.where(max_rgb == 0, 0, 1 - (min_rgb / max_rgb))  
    mean_saturation = np.mean(saturation_pixel) * 100  
    
    return [mean_brightness, edge_density, mean_saturation]


X_train, y_train = [], []
for path in cat_paths:
    X_train.append(extract_manual_features(path))
    y_train.append(0)
for path in dog_paths:
    X_train.append(extract_manual_features(path))
    y_train.append(1)

X_train = np.array(X_train)
y_train = np.array(y_train)


scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)


def get_test_data(test_cat_dir, test_dog_dir):
    """从测试目录收集图片路径和真实标签"""
    test_data = []  
    for filename in os.listdir(test_cat_dir):
        if filename.lower().endswith(IMAGE_EXTENSIONS):
            img_path = os.path.join(test_cat_dir, filename)
            test_data.append((img_path, 0))  
    for filename in os.listdir(test_dog_dir):
        if filename.lower().endswith(IMAGE_EXTENSIONS):
            img_path = os.path.join(test_dog_dir, filename)
            test_data.append((img_path, 1)) 
    return test_data


test_data = get_test_data(TEST_CAT_DIR, TEST_DOG_DIR)



print(f"\n=== KNN批量测试结果（共{len(test_data)}张测试图）===")
for k in [5, 7, 9]:
    knn = KNeighborsClassifier(n_neighbors=k)
    knn.fit(X_train, y_train)
    
    correct = 0
    total = len(test_data)
    print(f"\n--- K={k} 详细结果 ---")
    print("-" * 80)
    print(f"{'测试文件名':<20} {'预测结果':<10} {'真实结果':<10} {'是否正确'}")
    print("-" * 80)
    
    for img_path, true_label in test_data:
        test_feat = extract_manual_features(img_path)
        test_feat_norm = scaler.transform([test_feat])  
            

        pred_label = knn.predict(test_feat_norm)[0]
            

        pred_str = "猫" if pred_label == 0 else "狗"
        true_str = "猫" if true_label == 0 else "狗"
        is_correct = "是" if pred_label == true_label else "否"
            

        if is_correct == "是":
            correct += 1

        filename = os.path.basename(img_path)
        print(f"{filename:<20} {pred_str:<10} {true_str:<10} {is_correct}")

    accuracy = correct / total * 100
    print("-" * 80)
    print(f"K={k} 总结：共测试{total}张图，正确{correct}张，准确率：{accuracy:.2f}%")

print("\n=== 测试总结 ===")

accuracies = []
for k in [5, 7, 9]:
    knn = KNeighborsClassifier(n_neighbors=k)
    knn.fit(X_train, y_train)
    correct = 0
    for img_path, true_label in test_data:
        try:
            test_feat = extract_manual_features(img_path)
            test_feat_norm = scaler.transform([test_feat])
            pred_label = knn.predict(test_feat_norm)[0]
            if pred_label == true_label:
                correct += 1
        except:
            continue
    acc = correct / len([x for x in test_data if os.path.exists(x[0])]) * 100  
    accuracies.append((k, acc))


best_k, best_acc = max(accuracies, key=lambda x: x[1])
print(f"在当前测试集下，最优K值为 {best_k}，对应准确率：{best_acc:.2f}%")